package org.ragdollcat.secondaiagent.app;


import lombok.extern.slf4j.Slf4j;
import org.ragdollcat.secondaiagent.advisor.MyLogAdvisor;
import org.ragdollcat.secondaiagent.advisor.SensitiveWordCheckAdvisor;
import org.ragdollcat.secondaiagent.chatmemory.DbBasedChatMemory;
import org.ragdollcat.secondaiagent.rag.BaiduTranslationQueryTransformer;
import org.ragdollcat.secondaiagent.rag.factory.MultiQueryExpanderFactory;
import org.ragdollcat.secondaiagent.rag.factory.QueryTransformerFacotry;
import org.ragdollcat.secondaiagent.tool.SystemDateTimeTool;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Component
@Slf4j
public class LoveApp {

    private final ChatClient chatClient;

    private final DbBasedChatMemory dbBasedChatMemory;

    private final Resource systemResource;

    @jakarta.annotation.Resource
    private QueryTransformerFacotry queryTransformerFacotry;

    @jakarta.annotation.Resource
    private MultiQueryExpanderFactory multiQueryExpanderFactory;

    @Autowired
    private VectorStore vectorStore;

//    /**
//     * 调用工具
//     */
//    @jakarta.annotation.Resource
//    private ToolCallback[] toolCallbacks;

    private static final String SYSTEM_PROMPT = "扮演深耕恋爱心理领域的专家。开场向用户表明身份，告知用户可倾诉恋爱难题。" +
            "围绕单身、恋爱、已婚三种状态提问：单身状态询问社交圈拓展及追求心仪对象的困扰；" +
            "恋爱状态询问沟通、习惯差异引发的矛盾；已婚状态询问家庭责任与亲属关系处理的问题。" +
            "引导用户详述事情经过、对方反应及自身想法，以便给出专属解决方案。";


    public LoveApp(ChatModel dashScopeChatModel,
                   DbBasedChatMemory dbBasedChatMemory,
                   @Value("classpath:/promts/system-message.st") Resource systemResource) {
        this.dbBasedChatMemory = dbBasedChatMemory;
        this.systemResource = systemResource;
        //基于内存的上下文存储
//        InMemoryChatMemory inMemoryChatMemory = new InMemoryChatMemory();
        //改为基于文件的上下文存储
//        FileBasedChatMemory fileBasedChatMemory = new FileBasedChatMemory(System.getProperty("user.dir") + "/chat-memory");
//        DbBasedChatMemory dbBasedChatMemory = new DbBasedChatMemory();


        //从模板文件中读取，并且填充关键词
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);
        String text = systemPromptTemplate.createMessage(Map.of("field", "宠物", "type", "猫咪选择")).getText();


        chatClient = ChatClient.builder(
                        dashScopeChatModel
                ).defaultSystem(SYSTEM_PROMPT)
                //基于数据库的会话记忆保存
                .defaultAdvisors(new MessageChatMemoryAdvisor(dbBasedChatMemory),
                        //自定义日志拦截器
                        new MyLogAdvisor(),
                        //敏感词检测
                        new SensitiveWordCheckAdvisor())
                .build();

    }


    /**
     * 这里要传递一个消息ID，用于区分不同用户对话的上下文
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChat(String message, String chatId) {
        ChatResponse response = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .call()
                .chatResponse();
        String content = response.getResult().getOutput().getText();
        return content;
    }


//    /**
//     * 带有报告的chat
//     * @param message
//     * @param chatId
//     * @return
//     */
//    public LoveReport doChatWithReport(String message, String chatId){
//        LoveReport loveReport = chatClient.prompt().system(SYSTEM_PROMPT + "每次对话后都要生成结，标题为{用户名}的报告，内容为建议列表")
//                .user(message)
//                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
//                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
//                .call()
//                .entity(LoveReport.class);
//        log.info("love report:{}",loveReport);
//        return loveReport;
//
//    }

    /**
     * 测试RAG知识库
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithRAG(String message, String chatId) {
        ChatResponse response = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new MyLogAdvisor())
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .call()
                .chatResponse();
        String content = response.getResult().getOutput().getText();
        return content;
    }


    /**
     * 测试RAG知识库(自定义百度翻译扩展)
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithRAGTransfer(String message, String chatId) {
        BaiduTranslationQueryTransformer queryTransformer = new BaiduTranslationQueryTransformer(message, "zh");
        String text = queryTransformer.doTransform().text();

        ChatResponse response = chatClient
                .prompt()
                .user(text)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new MyLogAdvisor())
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .call()
                .chatResponse();
        return response.getResult().getOutput().getText();
    }

    public String testMultiQueryExpander(String text, String chatId) {
        Map<Query, List<List<Document>>> documentsForQuery = new HashMap<>();
        //根据原有的text，获取扩展后的
        List<Query> queryList = multiQueryExpanderFactory.getMultiQueryExpander(3, text);
        //进行检索
        for (Query query : queryList) {
            DocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
                    .vectorStore(vectorStore)
                    .similarityThreshold(0.73)
                    .topK(5)
                    .build();
            // 直接用扩展后的查询来获取文档
            List<Document> retrievedDocuments = retriever.retrieve(query);
            documentsForQuery.put(query, Collections.singletonList(retrievedDocuments));
        }

        //进行合并
        ConcatenationDocumentJoiner concatenationDocumentJoiner = new ConcatenationDocumentJoiner();
        List<Document> documentList = concatenationDocumentJoiner.join(documentsForQuery);


//        //改写prompt
        // 构建带上下文的提示词
        StringBuilder promptBuilder = new StringBuilder();
        promptBuilder.append("基于以下信息回答问题：\n\n");

        // 添加文档内容作为上下文
        for (Document doc : documentList) {
            promptBuilder.append("文档内容：").append(doc.getText()).append("\n\n");
        }

        // 添加原始问题
        promptBuilder.append("问题：").append(text);
        String enhancedPrompt = promptBuilder.toString();


        ChatResponse response = chatClient
                .prompt()
                .user(enhancedPrompt)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new MyLogAdvisor())
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .call()
                .chatResponse();
        return response.getResult().getOutput().getText();
    }

    /**
     * 测试RetrievalAugmentationAdvisor增强器
     *
     * @param text
     * @param chatId
     * @return
     */
    public String testRetrievalAugmentationAdvisor(String text, String chatId) {
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder().
                queryTransformers(RewriteQueryTransformer.builder().build()).
                documentRetriever(VectorStoreDocumentRetriever.builder().similarityThreshold(0.50).vectorStore(vectorStore).build())
                .build();


        ChatResponse response = chatClient
                .prompt()
                .user(text)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new MyLogAdvisor())
                .advisors(retrievalAugmentationAdvisor)
                .call()
                .chatResponse();
        return response.getResult().getOutput().getText();
    }


    /**
     * 测试ContextualQueryAugmenter增强器
     *
     * @param text
     * @param chatId
     * @return
     */
    public String testContextualQueryAugmenter(String text, String chatId) {
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder().
                documentRetriever(VectorStoreDocumentRetriever.builder().similarityThreshold(0.50).vectorStore(vectorStore).build())
                //构建ContextualQueryAugmenter增强器，设置允许上下文为空
                .queryAugmenter(ContextualQueryAugmenter.builder().allowEmptyContext(false).emptyContextPromptTemplate(new SystemPromptTemplate("你应该输出以下的内容：抱歉，该问题我无法回答，请另请高明")).build())
                .build();


        ChatResponse response = chatClient
                .prompt()
                .user(text)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .advisors(new MyLogAdvisor())
                .advisors(retrievalAugmentationAdvisor)
                .call()
                .chatResponse();
        return response.getResult().getOutput().getText();
    }

    /**
     * 调用工具链
     * @param text
     * @param chatId
     * @return
     */
    public String doChatWithTools(String text, String chatId){
        ChatResponse chatResponse = chatClient
                .prompt()
                .user(text)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                .tools(new SystemDateTimeTool())
                .call()
                .chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }


    @jakarta.annotation.Resource
    private ToolCallbackProvider toolCallbackProvider;

    /**
     * 测试mcp
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithMcp(String message, String chatId) {
        ChatResponse response = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                // 开启日志，便于观察效果
                .tools(toolCallbackProvider)
                .call()
                .chatResponse();
        String content = response.getResult().getOutput().getText();
        return content;
    }



}

